import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np


# 定义二维二次损失函数（最小值在原点）
def loss_fn(var):
    return var[0] ** 2 + 0.5 * var[1] ** 2  # x² + 0.5y²


# 初始化参数（初始位置设为(3,4)）
initial_var = tf.Variable([3.0, 4.0])

# 定义要对比的优化器（使用Keras优化器接口）
optimizers = {
    'SGD': tf.keras.optimizers.SGD(learning_rate=0.1),
    'Momentum': tf.keras.optimizers.SGD(learning_rate=0.1, momentum=0.9),
    'RMSprop': tf.keras.optimizers.RMSprop(learning_rate=0.1),
    'Adam': tf.keras.optimizers.Adam(learning_rate=0.1)
}

# 记录各优化器的参数更新轨迹
trajectories = {}
for name, optimizer in optimizers.items():
    # 每个优化器独立初始化参数
    var = tf.Variable(initial_var.numpy())
    path = [var.numpy()]  # 记录初始位置

    # 执行50次优化迭代
    for _ in range(50):
        optimizer.minimize(lambda: loss_fn(var), [var])  # Keras风格的最小化接口
        path.append(var.numpy())

    trajectories[name] = np.array(path)

# 可视化部分：绘制损失函数等高线和优化轨迹
x = np.linspace(-4, 4, 100)
y = np.linspace(-4, 4, 100)
X, Y = np.meshgrid(x, y)
Z = X ** 2 + 0.5 * Y ** 2  # 计算损失函数值

plt.figure(figsize=(10, 8))
plt.contour(X, Y, Z, levels=20, cmap='viridis', alpha=0.7)  # 绘制损失函数等高线

# 绘制各优化器的参数更新路径
for name, path in trajectories.items():
    plt.plot(path[:, 0], path[:, 1], marker='o', linestyle='-',
             label=name, linewidth=2, markersize=5)

plt.xlabel('Parameter x')
plt.ylabel('Parameter y')
plt.title('Optimizer Trajectory Comparison (Quadratic Loss)')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()